Fix: add missing before and after callback events for encode response #407
Conversation
bhimrazy
requested review from
lantiga,
aniketmaurya,
ethanwharris,
Andrei-Aksionov and
Borda
as code owners
for more information, see https://pre-commit.ci
bhimrazy
changed the title
| @@ -70,11 +70,14 @@ def run_streaming_loop( | |||
| ) | |||
| callback_runner.trigger_event(EventTypes.AFTER_PREDICT, lit_api=lit_api) | |||
|
|
|||
| callback_runner.trigger_event(EventTypes.BEFORE_ENCODE_RESPONSE, lit_api=lit_api) | |||
There was a problem hiding this comment.
this one is a bit tricky since y_enc_gen is a generator and the values are actually not materialized till it goes through these lines. so can't add it here.
There was a problem hiding this comment.
maybe we can show a warning to users when this callback hook is used with streaming.
There was a problem hiding this comment.
Sure, @aniketmaurya.
It seems that even the predict hook does not execute as expected since predict is also a generator. Additionally, it appears that all of the last four events are executed before the predict function runs.
Script I used to test
from litserve import Callback, LitAPI, LitServer
class SimpleStreamAPI(LitAPI):
def setup(self, device) -> None:
self.model = lambda x: x
def decode_request(self, request):
return request["input"]
def predict(self, x):
print(f"Predicting {x}")
for i in range(4):
yield self.model(i)
def encode_response(self, output):
for out in output:
yield {"output": out}
class LogCallback(Callback):
def on_before_predict(self, *args, **kwargs):
print("Before prediction")
def on_after_predict(self, *args, **kwargs):
print("After prediction")
def on_before_encode_response(self, *args, **kwargs):
print("Before encoding response")
def on_after_encode_response(self, *args, **kwargs):
print("After encoding response")
if __name__ == "__main__":
api = SimpleStreamAPI()
server = LitServer(api, stream=True, callbacks=[LogCallback()])
server.run(port=8000)One potential approach could be to move the "after event" part to execute after the generator has completed. However, this might not align semantically(in code) with the idea of executing it immediately after the respective functions.
or Alternatively, as you suggested, we could simply display a warning to users when this callback hook is used with streaming.
Looking forward to your thoughts on this 😊.
What does this PR do?
Adds missing before and after callback events for encode response in the streaming loop
Before submitting
PR review
Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in GitHub issues there's a high chance it will not be merged.
Did you have fun?
Make sure you had fun coding 🙃